package com.mdframe.ai.agent.config;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import com.mdframe.ai.agent.tools.GetTableSchemaFunction;
import com.mdframe.ai.agent.tools.MySQLSchemaService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Description;

import java.util.Map;
import java.util.function.Function;

/**
 * @date 2024/10/18
 */
@Slf4j
@Configuration
public class AssistantTools {
    
    
    @Autowired
    private MySQLSchemaService mySQLSchemaService;
    
    @Bean
    @Description("根据数据库表生成代码")
    public Function<GetTableSchemaFunction.Request, GetTableSchemaFunction.Response> getTableSchema() {
        return request -> {
            if (StrUtil.isNotBlank(request.tableName())) {
                Map<String, String> schemaMap = new MySQLSchemaService().getTableSchema(request.tableName());
                return new GetTableSchemaFunction.Response("columns:" + schemaMap.toString());
            }
            return new GetTableSchemaFunction.Response("请告诉我要查询的表名");
        };
    }
    
    @Bean("ToFile")
    @Description("生成之后的java代码需要保存到本地的文件中")
    public Function<SaveFileRequest, String> saveCodeToFile() {
        return request -> {
            log.info("java代码：{}", request);
            
            if (request != null && StrUtil.isNotBlank(request.code())) {
                // 保存到本地文件
                FileUtil.writeBytes(request.code().getBytes(), "/var/tmp/" + request.filePath());
                return "保存成功";
            }
            
            return null;
        };
    }
    
    public record SaveFileRequest(String code, String filePath) {
        
        @Override
        public String code() {
            return code;
        }
        
        @Override
        public String filePath() {
            return filePath;
        }
    }
}
